import os
import cv2
import numpy as np
from glob import glob
from tqdm import tqdm
from osgeo import gdal, osr
def read_tif(input_file):
    # 打开原始图像
    src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    if src_ds is None:
        print('Unable to open', input_file)
    return src_ds
def gdal_to_opencv(src_ds):
    # 打开TIFF图像
    # src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    # if src_ds is None:
    #     print('无法打开文件:', input_file)
    #     return None

    # 获取图像的宽度和高度
    width = src_ds.RasterXSize
    height = src_ds.RasterYSize
    bands = src_ds.RasterCount

    # 根据波段数决定图像的颜色模式
    if bands == 1:  # 单波段图像（灰度）
        band = src_ds.GetRasterBand(1)
        data = band.ReadAsArray(0, 0, width, height)
        img = np.array(data, dtype=np.uint8)
    elif bands == 3:  # 三波段图像（BGR）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)

        img = np.dstack((b, g, r))  # 注意这里的顺序
        img = np.array(img, dtype=np.uint8)
        return img
    elif bands == 4:  # 四波段图像（通常为BGR+NIR或BGR+Alpha）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)
        nir_band = src_ds.GetRasterBand(4)  # 假设第四个波段是NIR

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)
        nir = nir_band.ReadAsArray(0, 0, width, height)

        # 创建BGR图像
        img_bgr = np.dstack((b, g, r))  # 注意这里的顺序
        img_bgr = np.array(img_bgr, dtype=np.uint8)

        # 如果需要，也可以处理NIR波段
        # img_nir = np.array(nir, dtype=np.uint8)

        # 返回BGR图像
        return img_bgr
    else:
        print("不支持的波段数:", bands)
        return None


def convert_tif_to_opencv(input_folder, output_folder):
    # 创建输出文件夹（如果不存在）
    os.makedirs(output_folder, exist_ok=True)

    # 获取输入文件夹中的所有.tif文件
    tif_files = glob(os.path.join(input_folder, "*.tif"))

    # 使用tqdm显示进度
    with tqdm(total=len(tif_files), desc="Converting TIF to OpenCV format", unit="file") as pbar:
        for tif_file in tif_files:
            # 读取TIFF文件
            src_ds = read_tif(tif_file)
            if src_ds is None:
                continue

            # 转换为OpenCV格式
            img = gdal_to_opencv(src_ds)
            if img is None:
                continue

            # 构建输出文件路径
            output_file = os.path.join(output_folder, os.path.basename(tif_file).replace('.tif', '.jpg'))

            # 保存图像
            cv2.imwrite(output_file, img)

            # 更新进度条
            pbar.update(1)


if __name__ == "__main__":
    input_folder = r"E:\satLocate\fusai\basemap\tif"  # 输入文件夹路径
    output_folder = r"E:\satLocate\fusai\basemap\raster"  # 输出文件夹路径
    convert_tif_to_opencv(input_folder, output_folder)